package gomoku_qtable

import (
	"bindolabs/gomoku/modules/gomoku_env"
	"bindolabs/gomoku/modules/termbox_utils"
	"fmt"
	"time"

	"github.com/nsf/termbox-go"
)

func (agent *QLearningAgent) Play(filename string) error {
	// 这个函数与玩家play, 并打印每一步
	// Load Q-table from file
	if filename != "" {
		err := agent.LoadQTable(filename)
		if err != nil {
			return fmt.Errorf("error loading Q-table: %v", err)
		}
	}

	state := agent.env.Reset()
	agent.PlayWithHuman = true

	nextState := state
	done := false
	var err error
	var winner gomoku_env.PlayerType

	var playerX, playerY int

	for !done {
		// Agent's turn
		stateStr := stateToString(state, gomoku_env.Player1)
		actions := getPossibleActions(state)
		action := agent.ChooseAction(stateStr, actions)
		x, y := actionToCoordinates(action)

		nextState, done, winner, err = agent.env.Step(x, y)
		if err != nil {
			return fmt.Errorf("error during agent's action: %v", err)
		}
		agent.env.ShowBoard()
		agent.Update(state, nextState, action, done, winner)
		state = nextState
		// Check if the game is finished after agent's move
		if done {
			switch winner {
			case gomoku_env.Player1:
				{
					fmt.Println("Agent wins!")
				}
			case gomoku_env.Player2:
				{
					fmt.Println("You win!  It's a draw!")
				}
			default:
				{
					fmt.Println("平局")
				}
			}
			filename := fmt.Sprintf("playWithHuman_%v", gomoku_env.BoardSize)
			agent.SaveQTable(filename)

			state = agent.env.Reset()
			done = false
		}

		// Player's turn

		// Player's turn
	PlayerTurn:
		for {
			fmt.Println("Your move: use WASD or arrow keys to move, Enter to place:")
			fmt.Printf("Current position: (%d, %d)\n", playerX, playerY)
			input := termbox_utils.PollKeyboardEvent()

			switch input {
			case termbox.KeyArrowDown, termbox.KeyCtrlS:
				if playerX < gomoku_env.BoardSize-1 {
					playerX++
				}
			case termbox.KeyArrowUp, termbox.KeyCtrlW:
				if playerX > 0 {
					playerX--
				}
			case termbox.KeyArrowLeft, termbox.KeyCtrlA:
				if playerY > 0 {
					playerY--
				}
			case termbox.KeyArrowRight, termbox.KeyCtrlD:
				if playerY < gomoku_env.BoardSize-1 {
					playerY++
				}
			case termbox.KeyEnter:
				if state[playerX][playerY] == gomoku_env.Empty {
					nextState, done, winner, err = agent.env.Step(playerX, playerY)
					if err != nil {
						fmt.Printf("error during player's action: %v", err)
						continue
					}
					// agent.env.ShowBoard()
					// agent.Update(state, action, done, winner)
					break PlayerTurn // Exit inner loop after player places a piece
				} else {
					fmt.Println("Invalid move, try again.")
					continue
				}
			case termbox.KeyEsc:
				{
					fmt.Println("pressed esc exit.")
					return nil
				}
			}
			agent.env.ShowBoardWithCursor(playerX, playerY)
		}
		agent.Update(state, nextState, action, done, winner)
		state = nextState
		termbox_utils.ClearScreen()
		agent.env.ShowBoard()
		termbox.Flush()

		// Check if the game is finished after player's move
		if done {
			if winner == gomoku_env.Player2 {
				fmt.Println("You win!")
			} else {
				fmt.Println("It's a draw!")
			}

			filename := fmt.Sprintf("playWithHuman_%v", gomoku_env.BoardSize)
			agent.SaveQTable(filename)
			state = agent.env.Reset()
			done = false

			// return agent.SaveQTable(filename)
		}

	}

	return nil
}

func Run() {
	err := termbox.Init()
	if err != nil {
		panic(err.Error())
	}
	defer termbox.Close()

	env := gomoku_env.NewGomokuEnv()
	agent := NewQLearningAgent(env)

	if err := agent.Load(); err != nil {
		println(err.Error())
	}
	menu := []string{
		"0 exit",
		"1 play again",
		"2 choose qtable",
		"3 train qtable",
		"4 reset agent",
		"5 save agent",
		"6 show qtable",
		"7 menu watch agent",
	}

	done := false
	choose := 7

	for !done {
		switch choose {
		case 1:
			{
				if err := agent.Play(""); err != nil {
					println("game is end." + err.Error())
					panic(err.Error())
				}
			}
		case 2:
			{
				if err := agent.Load(); err != nil {
					println(err.Error())
					panic(err.Error())
				}
			}
		case 3:
			{
				agent.Train()
			}
		case 4:
			agent = NewQLearningAgent(env)
		case 5:
			var filename string
			fmt.Printf("\ninput filename:")
			fmt.Scanf("%s", &filename)
			filename = fmt.Sprintf("%v_%v", filename, gomoku_env.BoardSize)
			if err := agent.SaveQTable(filename); err != nil {
				println("save failed! due to err:", err.Error())
			} else {
				println("save ok. file:", filename)
			}
		case 6:
			fmt.Printf("qtable:%v", agent.qTable)
		case 7:
			agent.Watch()
		case 0:
			done = true
		default:
			println("esc for quit")
		}
		choose, err = termbox_utils.DisplayMenu(menu)
		if err != nil {
			println(err.Error())
			done = true
		}

	}
}

func (agent *QLearningAgent) Watch() {
	// watch agent play on itself
	state := agent.env.Reset()
	nextState := state
	done := false
	var err error
	var winner gomoku_env.PlayerType

	options := []string{
		"0 next ",
		"1 watch a play",
		"2 watch 5 plays",
		"3 vs with human.",
		"4 training.",
		"esc for quit or back to main menu.",
	}
	watch := 0
	op := 0
Run:
	for {
		stateStr := stateToString(state, agent.env.Player)
		actions := getPossibleActions(state)
		action := agent.ChooseAction(stateStr, actions)
		x, y := actionToCoordinates(action)
		nextState, done, winner, err = agent.env.Step(x, y)
		if err != nil {
			fmt.Printf("error during agent's action: %v", err)
			break Run
		}
		agent.env.ShowBoard()
		agent.Update(state, nextState, action, done, winner)
		state = nextState

		if watch > 0 {
			time.Sleep(time.Millisecond * 500)
		}

		if done {
			state = agent.env.Reset()
		}
		if done && watch > 0 {
			watch--
		}
		if watch == 0 {
			op, err = termbox_utils.DisplayMenu(options)
		}
		if err != nil {
			fmt.Printf("exit due to err:%v", err)
			return
		}
		switch op {
		case 1:
			watch = 1
		case 2:
			watch = 5
		case 3:
			agent.Play("")
		case 4:
			agent.Train()
		}
		op = 999
	}
}
